
#include <sys/types.h>

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <sys/time.h>
#include <sys/param.h>
#include <sys/queue.h>
#include <sys/tree.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <err.h>
#include <unistd.h>
#include <assert.h>
#include <time.h>

#include <event.h>
#include <evhttp.h>

#include "amproxy.gen.h"
#include "status.h"
#include "utils.h"
#include "log.h"

extern int debug;

struct stats statistics;

static struct pattern_obj good_patterns;
static struct pattern_obj bad_patterns;

static int status_patterns(struct pattern_obj *data, struct evbuffer *databuf);
static void site_print_analysis(struct evbuffer *databuf, struct site *site);
static void inform_cache_notfound(struct evhttp_request *request,
    const char *url);

/* structure where we keep track of sites */

#ifndef MIN
#define MIN(x, y) ((x) < (y) ? (x) : (y))
#endif

int
site_compare(struct site *a, struct site *b)
{
	static char atmp[HTTP_MAX_URL], btmp[HTTP_MAX_URL];
	char *a_url, *b_url;
	char *a_slash, *b_slash;
	int host_len, res;
	assert(strlen(a->url) >= sizeof(HTTP_PREFIX));
	assert(strlen(b->url) >= sizeof(HTTP_PREFIX));

	a_url = a->url;
	b_url = b->url;
	a_slash = strchr(a_url + sizeof(HTTP_PREFIX), '/');
	b_slash = strchr(b_url + sizeof(HTTP_PREFIX), '/');
	if (a_slash == NULL) {
		snprintf(atmp, sizeof(atmp), "%s/", a_url);
		a_url = atmp;
		a_slash = strchr(a_url + sizeof(HTTP_PREFIX), '/');
		assert(a_slash != NULL);
	}

	if (b_slash == NULL) {
		snprintf(btmp, sizeof(btmp), "%s/", b_url);
		b_url = btmp;
		b_slash = strchr(b_url + sizeof(HTTP_PREFIX), '/');
		assert(b_slash != NULL);
	}

	host_len = MIN((int)(a_slash - a_url), (int)(b_slash - b_url));
	res = strncasecmp(a_url, b_url, host_len);
	if (res)
		return (res);

	return strcmp(a_slash, b_slash);
}

//建立site_tree类型的分裂树，并声明一个对象root
static SPLAY_HEAD(site_tree, site) root;

SPLAY_PROTOTYPE(site_tree, site, node, site_compare);
SPLAY_GENERATE(site_tree, site, node, site_compare);

static int
find_url_in_patterns(struct patternq* head, const char *url)
{
	struct pattern *entry;
	TAILQ_FOREACH(entry, head, next) {
		if (match_url(url, entry->pattern_host, entry->pattern_uri))
			return (1);
	}

	return (0);
}

int 
site_same_as_parent(struct site *site)
{
	static char parent_host[1024];
	struct site *parent = site->parent;
	char *host, *uri;
	u_short port;

	if (parent == NULL)
		return (1);

	while (parent) {
		if (parent->parent == NULL)
			break;
		parent = parent->parent;
	}

	if (http_hostportfile(parent->url, &host, &port, &uri) == -1)
		return (0);
	strlcpy(parent_host, host, sizeof(parent_host));

	return (match_url(site->url, parent_host, NULL));
}

enum DANGER_TYPES
site_child_danger(struct site *site)
{
	enum DANGER_TYPES danger = UNKNOWN;

	if (site_same_as_parent(site))
		danger = HARMLESS;

	if (find_url_in_patterns(&good_patterns.head, site->url))
		danger = HARMLESS;
	if (find_url_in_patterns(&bad_patterns.head, site->url))
		danger = DANGEROUS;

	return (danger);
}

enum DANGER_TYPES
site_recurse_danger(struct site* site)
{
	enum DANGER_TYPES danger = site->danger;
	struct site *child;

	TAILQ_FOREACH(child, &site->children, next) {
		enum DANGER_TYPES cur = site_recurse_danger(child);
		if (cur > danger)
			danger = cur;
	}

	return (danger);
}

//分析页面的顶层调用
enum DANGER_TYPES
site_analyze_danger(struct site *site)
{
	enum DANGER_TYPES danger = HARMLESS;
	struct site *child;

	if (site->parent != NULL)
		return site_child_danger(site);

	if (site->html_size == 0 && TAILQ_FIRST(&site->children) == NULL)
		return (UNKNOWN);

	TAILQ_FOREACH(child, &site->children, next) {
		enum DANGER_TYPES cur = site_recurse_danger(child);
		if (cur > danger)
			danger = cur;
	}

	/* find the highest danger of children for the root */

	return (danger);
}

//统计恶意内容的总数
int
site_count_dangerous(struct site *site)
{
	struct site *child;
	int total = 0;

	TAILQ_FOREACH(child, &site->children, next) {
		total += site_count_dangerous(child);
	}

	if (site->danger == DANGEROUS)
		total += 1;

	return (total);
}

//分析页面
void
site_complete(int fd, short what, void *arg)
{
	struct site *site = arg;
	struct timeval tv;

	site->danger = site_analyze_danger(site);

	gettimeofday(&tv, NULL);
	timersub(&tv, &site->tv_change, &tv);
	if (tv.tv_sec >= IDLE_TIME) {
		site->flags |= ANALYSIS_COMPLETE;
	} else {
		timerclear(&tv);
		tv.tv_sec = 1;
		evtimer_add(&site->ev_complete, &tv);
	}
}

void
site_expire(int fd, short what, void *arg)
{
	struct site *site = arg;

	fprintf(stderr, "Expiring state for %s\n", site->url);
	site_free(site);
}

void
site_change_time(struct site *parent, struct timeval *tv)
{
	while (parent != NULL) {
		/* only expire from the top */
		if (parent->parent == NULL) {
			struct timeval tv_timeout;
			/* update the expiration time */
			timerclear(&tv_timeout);
			tv_timeout.tv_sec = STATE_EXPIRATION_TIME;

			evtimer_add(&parent->ev_timeout, &tv_timeout);
		}

		parent->tv_change = *tv;
		parent = parent->parent;
	}
}

//在分裂树中查找站点
struct site *
site_find(const char *url)
{
	struct site tmp;

	tmp.url = (char *)url;
	return (SPLAY_FIND(site_tree, &root, &tmp));
}

//发出新站点访问请求
struct site *
site_new(const char *url, const char *parent_url)
{
	struct site *res, tmp, *parent = NULL;
	struct timeval tv;

	tmp.url = (char *)url;
	if ((res = SPLAY_FIND(site_tree, &root, &tmp)) != NULL) {
		/* we already got a match - what now? */
		goto done;
	}

	if (parent_url != NULL) {
		tmp.url = (char *)parent_url;
		parent = SPLAY_FIND(site_tree, &root, &tmp);

		/* nobody should be able to fake a request */
		if (parent == NULL)
			return (NULL);
	}

	if ((res = calloc(1, sizeof(struct site))) == NULL)
		err(1, "calloc");

	TAILQ_INIT(&res->children);
	if (parent != NULL) {
		res->parent = parent;
		TAILQ_INSERT_TAIL(&parent->children, res, next);
	}

	if ((res->url = strdup(url)) == NULL)
		err(1, "strdup");

	res->danger = site_analyze_danger(res);
	SPLAY_INSERT(site_tree, &root, res);

	evtimer_set(&res->ev_timeout, site_expire, res);
	evtimer_set(&res->ev_complete, site_complete, res);

done:
	/* update the last time a tree was updated */
	gettimeofday(&tv, NULL);
	site_change_time(res, &tv);

	return (res);
}

void
site_free(struct site *site)
{
	struct site *child;

	SPLAY_REMOVE(site_tree, &root, site);

	event_del(&site->ev_timeout);
	event_del(&site->ev_complete);

	while ((child = TAILQ_FIRST(&site->children)) != NULL) {
		TAILQ_REMOVE(&site->children, child, next);
		child->parent = NULL;
		site_free(child);
	}

	if (site->parent) {
		TAILQ_REMOVE(&site->parent->children, site, next);
	}

	if (site->html_data != NULL)
		free(site->html_data);
	free(site->url);
	free(site);
}

#define HTML_PRINT(...) evbuffer_add_printf(databuf, __VA_ARGS__)

/* stores the data associated with this site */

void
site_cache_data(struct site *site, const struct evhttp_request *req)
{

	if (site->html_data != NULL)
		free(site->html_data);

	fprintf(stderr, "Caching %ld bytes for %s (%s)\n",
	    EVBUFFER_LENGTH(req->input_buffer),
	    site->url, danger_to_text(site->danger));

	site->html_size = EVBUFFER_LENGTH(req->input_buffer);
	site->html_data = malloc(site->html_size);
	if (site->html_data == NULL)
		err(1, "malloc");

	memcpy(site->html_data, EVBUFFER_DATA(req->input_buffer),
	    site->html_size);


}


const char *
danger_to_text(enum DANGER_TYPES danger)
{
	switch (danger) {
	case HARMLESS:
		return "harmless";
	case DANGEROUS:
		return "dangerous";
	case UNKNOWN:
	default:
		return "unknown";
	}
}

static void
site_print_children(struct evbuffer *databuf, struct site *site,
    enum DANGER_TYPES desired_level)
{
	struct site *child;
	TAILQ_FOREACH(child, &site->children, next) {
		if (child->danger != desired_level)
			continue;

		HTML_PRINT("<li>");
		site_print_analysis(databuf, child);
		HTML_PRINT("</li>");
	}
}

static void
site_print_analysis(struct evbuffer *databuf, struct site *site)
{
	char *uri_escaped = evhttp_encode_uri(site->url);
	char *html_escaped = evhttp_htmlescape(site->url);
	HTML_PRINT("<span class=%s>%s</span> ", 
	    danger_to_text(site->danger),
	    danger_to_text(site->danger));

	HTML_PRINT(
		"<a href=\"/cache/?url=%s\" target=\"_blank\">%s</a>",
		uri_escaped, html_escaped);

	free(uri_escaped);
	free(html_escaped);
	if (TAILQ_FIRST(&site->children) != NULL) {
		HTML_PRINT("<ul>");
		site_print_children(databuf, site, DANGEROUS);
		site_print_children(databuf, site, UNKNOWN);
		site_print_children(databuf, site, HARMLESS);
		HTML_PRINT("</ul>");
	}
}



void
css_server(struct evhttp_request *request, void *arg)
{
	struct evbuffer *databuf = evbuffer_new();
	assert(databuf != NULL);
	evhttp_add_header(request->output_headers, "Content-Type", "text/css");
	evbuffer_add(databuf, css_style, strlen(css_style));

	/* send along our data */
	evhttp_send_reply(request, HTTP_OK, "OK", databuf);
	evbuffer_free(databuf);
}

static void
print_blurb(struct evbuffer *databuf)
{
	HTML_PRINT("<div class=footer><hr />\n");
	HTML_PRINT("<center>"
		"Version %s "
	    "Copyright (c) 2007 <a href=http://www.lovetech.cn/>"
	    "yn</a>.  All Rights Reserved."
	    "</center>\n",
		VERSION
		);
	HTML_PRINT("</div>\n");
}

static void
print_footer(struct evbuffer *databuf)
{
	HTML_PRINT("</body></html>");
}

static void
print_header(struct evbuffer *databuf)
{

	HTML_PRINT(
		"<html><head><title>"
		"AmProxy"
		"</title></head>");
	HTML_PRINT("<link rel=stylesheet type=text/css href=/styles/css>\n");

	HTML_PRINT(
		"<body><div class=banner>\n"
		"<h1>AmProxy</h1>\n"
		"<ul>"
		"<li>Anti-Malware Proxy</li>\n"
		"</ul>\n"
		"</div>\n"
	    );
}

static void
print_form(struct evbuffer *databuf)
{
	//提交类似/?url=xxx的地址
		
	HTML_PRINT(
	"<p><center>\n"
	"<form name=\"input\" action=\"/\" method=\"get\">\n"
	"Web Address: <input type=\"text\" name=\"url\" size=100>\n"
	"<input type=\"submit\" value=\"Go\">\n"
	"</form></center>");
}


//主页
static void
main_server(struct evhttp_request *request, void *arg)
{
	struct evbuffer *databuf = evbuffer_new();
	assert(databuf != NULL);

	print_header(databuf);

	print_form(databuf);

	print_blurb(databuf);

	print_footer(databuf);

	/* send along our data */
	evhttp_send_reply(request, HTTP_OK, "OK", databuf);
	evbuffer_free(databuf);
}



//分析结果窗口
void
results_server(struct evhttp_request *request, void *arg)
{
	struct site *site, tmp;
	struct evkeyvalq args;
	struct evbuffer *databuf = evbuffer_new();
	const char *url = NULL;
	int done = 0;
	assert(databuf != NULL);

	TAILQ_INIT(&args);

	evhttp_parse_query(request->uri, &args);

	url = evhttp_find_header(&args, "url");

	if (url == NULL)
		goto fail;

	tmp.url = (char *)url;
	site = SPLAY_FIND(site_tree, &root, &tmp);

	HTML_PRINT(
		"<html><head><title>"
		"AmProxy: Results</title>\n");
	done = site == NULL || (site->flags & ANALYSIS_COMPLETE);
	if (!done)
		HTML_PRINT("<meta http-equiv=\"refresh\" content=\"2\">\n");

	HTML_PRINT("</head>");
	HTML_PRINT("<link rel=stylesheet type=text/css href=/styles/css>\n");

	HTML_PRINT("<body>");

	tmp.url = (char *)url;
	if (site != NULL) {
		char *url_escaped = evhttp_htmlescape(site->url);
		HTML_PRINT("%s %s found %d dangerous links.<p>\n",
		    done ? "<span class=harmless>Complete</span>" :
		    "<span class=unknown>Pending</span>",
		    url_escaped, site_count_dangerous(site));
		free(url_escaped);

		HTML_PRINT("<div class=analysis>\n");
		site_print_analysis(databuf, site);
		HTML_PRINT("</div>\n");
	}

	if (done) {
		HTML_PRINT("<p>The analysis of this URL is complete. ");
	} else {
		struct timeval tv;
		gettimeofday(&tv, NULL);
		timersub(&tv, &site->tv_change, &tv);
	
		HTML_PRINT("<p>Analysis is going to take %d more seconds.",
		    IDLE_TIME - tv.tv_sec);
	}
	
	HTML_PRINT("</body></html>");

	evhttp_clear_headers(&args);

	/* send along our data */
	evhttp_send_reply(request, HTTP_OK, "OK", databuf);
	evbuffer_free(databuf);
	return;

fail:
	evhttp_clear_headers(&args);
	evhttp_send_error(request, HTTP_BADREQUEST, "You must be kidding.");
}

void
query_server(struct evhttp_request *request, void *arg)
{
	struct evkeyvalq args;
	struct evbuffer *databuf = evbuffer_new();
	const char *url = NULL;
	assert(databuf != NULL);

	TAILQ_INIT(&args);

	evhttp_parse_query(request->uri, &args);

	url = evhttp_find_header(&args, "url");

	print_header(databuf);

	print_form(databuf);

	print_blurb(databuf);

	if (url == NULL || http_hostportfile(url, NULL, NULL, NULL) == -1) {
		HTML_PRINT("Did not receive a URL.  You loose.");
	} else {
		struct site *site = site_new(url, NULL);
		char *url_relative = evhttp_encode_uri(url);

		if (site == NULL)
			goto fail;
		HTML_PRINT(
			"<iframe src=\"/results/?url=%s\" "
			"width=100%% height=30%%>\n"
			"</iframe><p>\n",
			url_relative);
		free(url_relative);

	//	if ((site->flags & ANALYSIS_COMPLETE) == 0) {
			char *url_escaped = evhttp_htmlescape(url);
			HTML_PRINT(
				"<iframe src=\"%s\" "
				"width=100%% height=50%% security=restricted>\n"
				"</iframe>\n",
				url_escaped);
			free(url_escaped);

			/* start the completion timer */
			site_complete(-1, 0, site);
	//	}
	}

	print_footer(databuf);

	evhttp_clear_headers(&args);

	/* send along our data */
	evhttp_send_reply(request, HTTP_OK, "OK", databuf);
	evbuffer_free(databuf);
	return;

fail:
	evhttp_clear_headers(&args);
	evhttp_send_error(request, HTTP_SERVUNAVAIL, "Another eval is on");
}

void
cache_server(struct evhttp_request *request, void *arg)
{
	struct timeval tv;
	struct evkeyvalq args;
	struct evbuffer *databuf = evbuffer_new();
	struct site *site;
	const char *url = NULL;
	char *escaped;
	assert(databuf != NULL);

	TAILQ_INIT(&args);

	evhttp_parse_query(request->uri, &args);

	url = evhttp_find_header(&args, "url");
	site = site_find(url);

	if (site == NULL || site->html_data == NULL)
		goto fail;

	evhttp_clear_headers(&args);

	/* somebody showed interst in this page - let it not expire yet */
	gettimeofday(&tv, NULL);
	site_change_time(site, &tv);

	/* NUL terminate */
	evbuffer_add(databuf, site->html_data, site->html_size);
	evbuffer_add(databuf, "", 1);

	escaped = evhttp_htmlescape((char *)EVBUFFER_DATA(databuf));

	evbuffer_drain(databuf, -1);

	HTML_PRINT("<html><head><title>raw dump</title></head><body>");
	HTML_PRINT("<pre>%s</pre>", escaped);
	free(escaped);
	HTML_PRINT("</body></html>");

	/* send along our data */
	evhttp_send_reply(request, HTTP_OK, "OK", databuf);
	evbuffer_free(databuf);
	return;

fail:
	inform_cache_notfound(request, url);
	evhttp_clear_headers(&args);
}

void
amproxy_handle_request(struct evhttp_request *request, void *arg)
{
	char *host, *uri;
	u_short port;

	if (http_hostportfile(request->uri, &host, &port, &uri) == -1) {
		evhttp_send_error(request,
		    HTTP_BADREQUEST, "Illegal request.");
		return;
	}

	/*
	 * this is a little bit silly, we are not taking advantage of
	 * the http layer dispatch support.
	 */
	//根据amproxy后的地址跳转到不同的网页
	if (strcmp(uri, "/styles/css") == 0) {
		css_server(request, arg);
		return;
	} else if (strcmp(uri, "/") == 0) {
		main_server(request, arg);
		return;
	} else if (strncmp(uri, "/?", 2) == 0) {
		query_server(request, arg);
		return;
	} else if (strncmp(uri, "/results/", 9) == 0) {
		results_server(request, arg);
		return;
	} else if (strncmp(uri, "/cache/", 7) == 0) {
		cache_server(request, arg);
		return;
	}

	evhttp_send_error(request, HTTP_NOTFOUND, "Data not found.");
}




void
status_init(const char *whitelist, const char *blacklist)
{
	struct evbuffer *data;
	SPLAY_INIT(&root);

	good_patterns.location = whitelist;
	bad_patterns.location = blacklist;


	/* initial setup of the contents */
		data =read_data(whitelist) ;
		status_good_patterns(data);
		evbuffer_free(data);
	

		data =read_data(blacklist) ;
		status_bad_patterns(data);
		evbuffer_free(data);
	
}

static int
status_patterns(struct pattern_obj *data, struct evbuffer *databuf)
{
	char *line;
	int count = 0;

	while ((line = evbuffer_readline(databuf)) != NULL) {
		struct pattern *pattern;
		char *host = line, *uri;
		if (line[0] == '#' || !strlen(line))
			continue;

		uri = strchr(line, ' ');
		if (uri != NULL)
			*uri++ = '\0';

		pattern = malloc(sizeof(struct pattern));
		if (pattern == NULL)
			err(1, "malloc");
		pattern->pattern_host = strdup(host);
		if (uri != NULL)
			pattern->pattern_uri = strdup(uri);
		else
			pattern->pattern_uri = NULL;

		if (pattern->pattern_host == NULL ||
		    (uri != NULL && pattern->pattern_uri == NULL))
			err(1, "strdup");

		TAILQ_INSERT_TAIL(&data->head, pattern, next);
		if (debug)
			fprintf(stderr, "Adding pattern: %s/%s\n",
			    host, uri);
		count++;

		free(line);
	}
	data->count = count;
	gettimeofday(&data->tv_load, NULL);

	return (count);
}

void
status_good_patterns(struct evbuffer *data)
{
	int count;
	TAILQ_INIT(&good_patterns.head);
	count = status_patterns(&good_patterns, data);

	fprintf(stderr, "Added %d good patterns\n", count);
}

void
status_bad_patterns(struct evbuffer *data)
{
	int count;
	TAILQ_INIT(&bad_patterns.head);
	count = status_patterns(&bad_patterns, data);

	fprintf(stderr, "Added %d bad patterns\n", count);
}


/* error messaging */

static void
inform_cache_notfound(struct evhttp_request *request, const char *url)
{
	struct evbuffer *databuf = evbuffer_new();
	char *html_escaped = evhttp_htmlescape(url != NULL ? url : "<unknown>");
	assert(databuf != NULL);

	evbuffer_add_printf(databuf,
	    "<html><head><title>Cache Not Found</title></head>"
	    "<body><div style=\"border: solid 1px; padding: 2px; "
	    "width: 40%%; "
	    "background-color: #dcdcee; font-family: Verdana, Arial;\">"
	    "<h2>Cache Not Found</h2>\n"
	    "The URL %s that you requested could not be found in the cache. "
	    "It's possible that it could not be fetched from its web "
	    "server.</div></body></html>",
	    html_escaped);
	free(html_escaped);

	/* we cannot allow this request */
	evhttp_send_reply(request, HTTP_NOTFOUND, "Not Found", databuf);
	evbuffer_free(databuf);
}
